from typing import List, Dict, Any, Optional
from abc import ABC, abstractmethod
from pydantic import BaseModel


class Node(BaseModel):
    """图节点基类"""
    id: str
    type: str
    properties: Dict[str, Any]
    embedding: Optional[List[float]] = None


class Edge(BaseModel):
    """图边基类"""
    source_id: str
    target_id: str
    type: str
    properties: Dict[str, Any]


class SubGraph(BaseModel):
    """子图结构"""
    nodes: List[Node]
    edges: List[Edge]


class GraphSearchResult(BaseModel):
    """图搜索结果"""
    subgraph: SubGraph
    score: float
    source: str
    metadata: Dict[str, Any]


class BaseGraphStore(ABC):
    """图数据库基类"""
    
    @abstractmethod
    def initialize(self, config: Dict[str, Any]) -> None:
        """初始化图数据库连接"""
        pass
    
    @abstractmethod
    def add_node(self, node: Node) -> None:
        """添加节点"""
        pass
    
    @abstractmethod
    def add_edge(self, edge: Edge) -> None:
        """添加边"""
        pass
    
    @abstractmethod
    def get_node(self, node_id: str) -> Optional[Node]:
        """获取节点"""
        pass
    
    @abstractmethod
    def get_neighbors(self, node_id: str, edge_types: Optional[List[str]] = None) -> List[Node]:
        """获取邻居节点"""
        pass
    
    @abstractmethod
    def search_nodes(self, query: str, node_types: Optional[List[str]] = None) -> List[Node]:
        """搜索节点"""
        pass
    
    @abstractmethod
    def get_subgraph(self, node_ids: List[str], depth: int = 1) -> SubGraph:
        """获取子图"""
        pass


class BaseGraphEmbedder(ABC):
    """图嵌入基类"""
    
    @abstractmethod
    def initialize(self, config: Dict[str, Any]) -> None:
        """初始化嵌入模型"""
        pass
    
    @abstractmethod
    def embed_text(self, text: str) -> List[float]:
        """文本嵌入"""
        pass
    
    @abstractmethod
    def embed_node(self, node: Node) -> List[float]:
        """节点嵌入"""
        pass
    
    @abstractmethod
    def embed_subgraph(self, subgraph: SubGraph) -> List[float]:
        """子图嵌入"""
        pass


class BaseGraphSearcher(ABC):
    """图搜索基类"""
    
    @abstractmethod
    def initialize(self, config: Dict[str, Any]) -> None:
        """初始化搜索器"""
        pass
    
    @abstractmethod
    async def search(self, query: str, **kwargs) -> List[GraphSearchResult]:
        """执行搜索"""
        pass


class BaseGraphAugmenter(ABC):
    """图增强基类"""
    
    @abstractmethod
    def initialize(self, config: Dict[str, Any]) -> None:
        """初始化增强器"""
        pass
    
    @abstractmethod
    async def augment(self, query: str, subgraph: SubGraph) -> SubGraph:
        """增强子图"""
        pass


class BaseGraphReasoner(ABC):
    """图推理基类"""
    
    @abstractmethod
    def initialize(self, config: Dict[str, Any]) -> None:
        """初始化推理器"""
        pass
    
    @abstractmethod
    async def reason(self, query: str, subgraph: SubGraph) -> Dict[str, Any]:
        """执行推理"""
        pass 